import torch
import numpy as np
from models.inception_res_v1 import InceptionResnetV1

seq_conv_map_dict = {'conv2d_1a': 'Conv2d_1a_3x3',
                     'conv2d_2a': 'Conv2d_2a_3x3',
                     'conv2d_2b': 'Conv2d_2b_3x3',
                     'conv2d_3b': 'Conv2d_3b_1x1',
                     'conv2d_4a': 'Conv2d_4a_3x3',
                     'conv2d_4b': 'Conv2d_4b_3x3',
                     }

block_35_map_dict = {'branch0_1x1': 'Branch_0/Conv2d_1x1',
                     'branch1_0_1x1': 'Branch_1/Conv2d_0a_1x1',
                     'branch1_1_3x3': 'Branch_1/Conv2d_0b_3x3',
                     'branch2_0_1x1': 'Branch_2/Conv2d_0a_1x1',
                     'branch2_1_3x3': 'Branch_2/Conv2d_0b_3x3',
                     'branch2_2_3x3': 'Branch_2/Conv2d_0c_3x3',
                     'conv2d': 'Conv2d_1x1',
                     }

block_17_map_dict = {'branch0_1x1': 'Branch_0/Conv2d_1x1',
                     'branch1_0_1x1': 'Branch_1/Conv2d_0a_1x1',
                     'branch1_1_1x7': 'Branch_1/Conv2d_0b_1x7',
                     'branch1_2_7x1': 'Branch_1/Conv2d_0c_7x1',
                     'conv2d': 'Conv2d_1x1',
                     }

block_8_map_dict = {'branch0_1x1': 'Branch_0/Conv2d_1x1',
                    'branch1_0_1x1': 'Branch_1/Conv2d_0a_1x1',
                    'branch1_1_1x3': 'Branch_1/Conv2d_0b_1x3',
                    'branch1_2_3x1': 'Branch_1/Conv2d_0c_3x1',
                    'conv2d': 'Conv2d_1x1',
                    }
mixed6a_map_dict = {'branch0_3x3': 'Branch_0/Conv2d_1a_3x3',
                    'branch1_0_1x1': 'Branch_1/Conv2d_0a_1x1',
                    'branch1_1_3x3': 'Branch_1/Conv2d_0b_3x3',
                    'branch1_2_3x3': 'Branch_1/Conv2d_1a_3x3',
                    }

mixed7a_map_dict = {'branch0_0_1x1': 'Branch_0/Conv2d_0a_1x1',
                    'branch0_1_3x3': 'Branch_0/Conv2d_1a_3x3',
                    'branch1_0_1x1': 'Branch_1/Conv2d_0a_1x1',
                    'branch1_1_3x3': 'Branch_1/Conv2d_1a_3x3',
                    'branch2_0_1x1': 'Branch_2/Conv2d_0a_1x1',
                    'branch2_1_3x3': 'Branch_2/Conv2d_0b_3x3',
                    'branch2_2_3x3': 'Branch_2/Conv2d_1a_3x3',
                    }

conv_map_dict = {'conv.weight': 'weights',
                 'bn.bias': 'BatchNorm/beta',
                 'bn.running_mean': 'BatchNorm/moving_mean',
                 'bn.running_var': 'BatchNorm/moving_variance',
                 'weight': 'weights',
                 'bias': 'biases'}

fc_map_dict = {'fc1.weight':'bottleneck/weights',
               'fc1.bias': 'bottleneck/biases',
               'fc2.weight':'logits/weights'}

def tf_to_ptch(tf_weights,pt_weights):
    '''
    :param tf_weights:
    :param pt_weights:
    :return:
    '''
    converted_weights = {}

    # N,C,H,W in pytorch
    for item in pt_weights.keys():
        if item.split('.')[0] in ['conv2d_1a', 'conv2d_2a',
                                  'conv2d_2b', 'conv2d_3b',
                                  'conv2d_4a', 'conv2d_4b']:
            conv_name = item.split('.')[0]
            weights_name = '.'.join(item.split('.')[1:])
            if weights_name == 'bn.weight':
                tg_weights = np.ones(shape=pt_weights[item].size(), dtype=np.float32)
            else:
                tf_name = '%s/%s' % (seq_conv_map_dict[conv_name], conv_map_dict[weights_name])
                # tg_weights = np.empty(shape=tf_weights[tf_name])
                tg_weights = tf_weights[tf_name]
            if len(tg_weights.shape) >= 4:
                tg_weights = tg_weights.transpose(3, 2, 0, 1)
            tg_weights = torch.from_numpy(tg_weights)

        elif item.split('.')[0] == 'mixed_6a':
            branch_name = item.split('.')[1]
            weights_name = '.'.join(item.split('.')[2:])
            if weights_name == 'bn.weight':
                tg_weights = np.ones(shape=pt_weights[item].size(), dtype=np.float32)
            else:
                tf_name = 'Mixed_6a/%s/%s' % (mixed6a_map_dict[branch_name],
                                              conv_map_dict[weights_name])
                # tg_weights = np.empty(shape=tf_weights[tf_name])
                tg_weights = tf_weights[tf_name]
            if len(tg_weights.shape) >= 4:
                tg_weights = tg_weights.transpose(3, 2, 0, 1)
            tg_weights = torch.from_numpy(tg_weights)

        elif item.split('.')[0] == 'mixed_7a':
            branch_name = item.split('.')[1]
            weights_name = '.'.join(item.split('.')[2:])
            if weights_name == 'bn.weight':
                tg_weights = np.ones(shape=pt_weights[item].size(), dtype=np.float32)
            else:
                tf_name = 'Mixed_7a/%s/%s' % (mixed7a_map_dict[branch_name],
                                              conv_map_dict[weights_name])
                # tg_weights = np.empty(shape=tf_weights[tf_name])
                tg_weights = tf_weights[tf_name]
            if len(tg_weights.shape) >= 4:
                tg_weights = tg_weights.transpose(3, 2, 0, 1)
            tg_weights = torch.from_numpy(tg_weights)

        elif item.split('.')[0] == 'block8':
            branch_name = item.split('.')[1]
            weights_name = '.'.join(item.split('.')[2:])
            if weights_name == 'bn.weight':
                tg_weights = np.ones(shape=pt_weights[item].size(), dtype=np.float32)
            else:
                tf_name = 'Block8/%s/%s' % (block_8_map_dict[branch_name],
                                            conv_map_dict[weights_name])
                # tg_weights = np.empty(shape=tf_weights[tf_name])
                tg_weights = tf_weights[tf_name]
            if len(tg_weights.shape) >= 4:
                tg_weights = tg_weights.transpose(3, 2, 0, 1)

            tg_weights = torch.from_numpy(tg_weights)

        elif 'repeat_block35' in item:
            block_idx = int(item.split('.')[1]) + 1
            branch_name = item.split('.')[2]
            weights_name = '.'.join(item.split('.')[3:])
            if weights_name == 'bn.weight':
                tg_weights = np.ones(shape=pt_weights[item].size(), dtype=np.float32)
            else:
                tf_name = 'Repeat/block35_%s/%s/%s' % (block_idx,
                                                       block_35_map_dict[branch_name],
                                                       conv_map_dict[weights_name])
                # tg_weights = np.empty(shape=tf_weights[tf_name])
                tg_weights = tf_weights[tf_name]
            if len(tg_weights.shape) >= 4:
                tg_weights = tg_weights.transpose(3, 2, 0, 1)
            tg_weights = torch.from_numpy(tg_weights)

        elif 'repeat_block17' in item:
            block_idx = int(item.split('.')[1]) + 1
            branch_name = item.split('.')[2]
            weights_name = '.'.join(item.split('.')[3:])
            if weights_name == 'bn.weight':
                tg_weights = np.ones(shape=pt_weights[item].size(), dtype=np.float32)
            else:

                tf_name = 'Repeat_1/block17_%s/%s/%s' % (block_idx,
                                                         block_17_map_dict[branch_name],
                                                         conv_map_dict[weights_name])
                # tg_weights = np.empty(shape=tf_weights[tf_name])
                tg_weights = tf_weights[tf_name]
            if len(tg_weights.shape) >= 4:
                tg_weights = tg_weights.transpose(3, 2, 0, 1)
            tg_weights = torch.from_numpy(tg_weights)
        elif 'repeat_block8' in item:
            block_idx = int(item.split('.')[1]) + 1
            branch_name = item.split('.')[2]
            weights_name = '.'.join(item.split('.')[3:])
            if weights_name == 'bn.weight':
                tg_weights = np.ones(shape=pt_weights[item].size(), dtype=np.float32)
            else:

                tf_name = 'Repeat_2/block8_%s/%s/%s' % (block_idx,
                                                        block_8_map_dict[branch_name],
                                                        conv_map_dict[weights_name])
                # tg_weights = np.empty(shape=tf_weights[tf_name])
                tg_weights = tf_weights[tf_name]

            if len(tg_weights.shape) >= 4:
                tg_weights = tg_weights.transpose(3, 2, 0, 1)
            tg_weights = torch.from_numpy(tg_weights)
        elif 'fc' in item:
            tf_name = fc_map_dict[item]
            tg_weights = tf_weights[tf_name]
            if item in ['fc1.weight','fc2.weight']  :
                tg_weights = tg_weights.transpose(1,0)
            tg_weights = torch.from_numpy(tg_weights)
        else:
            tg_weights = pt_weights[item]
            print 'unmatched:', item, ',remain pytorch weights'
        if tg_weights.shape != pt_weights[item].size():
            print 'inconsistent weights size: "%s" and "%s"' % (item, tf_name),
            print pt_weights[item].size(),tg_weights.shape

        converted_weights[item] = tg_weights
    return converted_weights



if __name__ == '__main__':
    model = InceptionResnetV1(num_classes=75322, image_channel=1)
    model.eval()
    # print model

    x = torch.FloatTensor(10, 1, 112, 96)
    x = torch.autograd.Variable(x)

    y = model(x)

    pt_weights = model.state_dict()
    tf_weights = np.load('../trained_weights.npy')[()]

    # for item in pt_weights.keys():
    #     if 'fc' in item:
    #         print item,pt_weights[item].shape
    #
    # for item in tf_weights.keys():
    #     if 'bottleneck' in item:
    #         print item,tf_weights[item].shape

    converted_weights = tf_to_ptch(tf_weights, pt_weights)
    #
    model.load_state_dict(converted_weights)
